from lib.visual_detect import detect_sidewalk, detect_traffic_light, detect_yellow_line
from std_msgs.msg import String
import rospy
import cv2
import json
import os
import time

# 配置字典
config = {
    'camera_index': 0, # 摄像头索引，0 通常是默认摄像头
    'video_filename': f'real_{int(time.time())}.avi', # 输出视频文件的名称
    'video_codec': 'XVID', # 视频编码器
    'fps': 30.0, # 视频的帧率
    'frame_size': (640, 480), # 视频的分辨率
}

# 打开摄像头
video = cv2.VideoCapture(config['camera_index'])

# 检查摄像头是否成功打开
if not video.isOpened():
    print("无法打开摄像头")
    exit()

# 定义视频编码器和创建VideoWriter对象
fourcc = cv2.VideoWriter_fourcc(*config['video_codec'])
out = cv2.VideoWriter(config['video_filename'], fourcc, config['fps'], config['frame_size'])


rospy.init_node("visual_node", anonymous=True)
pub = rospy.Publisher("/visual", String, queue_size=2)
print("[VISUAL]")
# 保存视屏
try:
    while video.isOpened():
        ret, frame = video.read()
        if ret is True:
            frame = cv2.flip(frame, 1)
            frame = cv2.flip(frame, 0)
            out.write(frame)
            sidewalk_flag = detect_sidewalk(frame)
            yellow_line_flag = detect_yellow_line(frame)
            traffic_light_flag = detect_traffic_light(frame)

            res = {
                "sidewalk": sidewalk_flag,
                "yellow_line": yellow_line_flag,
                "traffic_light": traffic_light_flag,
            }

            if res["sidewalk"]or res["yellow_line"]or res["traffic_light"]:
                print(res)
                pub.publish(json.dumps(res))

            cv2.imshow("Frame", frame)
            ch = cv2.waitKey(1) & 0xFF
            if ch == ord("q"):
                break
        else:
            break
except Exception as e:
    print(e)
finally:
    video.release()
    out.release()
    cv2.destroyAllWindows()
    os.system('pause')
